{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "675904d8",
   "metadata": {},
   "source": [
    "# Tutorial 08: creating nullspace vector for pure Neumann problem with restrictions\n",
    "\n",
    "In this tutorial we solve the problem\n",
    "\n",
    "$$\\begin{cases}\n",
    "-\\Delta u = f, & \\text{in } \\Omega,\\\\\n",
    "\\nabla u \\cdot\\mathbf{n} = g, & \\text{on } \\partial\\Omega,\n",
    "\\end{cases}$$\n",
    "\n",
    "where $\\Omega$ is a ball in 2D. The forcing term $f$ and the Neumann data $g$ are such that they satisfy the condition\n",
    "$$\\int_\\Omega f \\; dx = - \\int_{\\partial\\Omega} g \\; ds $$\n",
    "which is a necessary condition for the existence of the solution to this pure Neumann problem, and which can be easily obtained multiplying the first equation in the system by the constant 1 and integrating by parts. Note that the solution is determined up to a constant.\n",
    "\n",
    "The domain $\\Omega$ is decomposed in $\\Omega = \\Omega_1 \\cup \\Omega_2$ with $\\Gamma$ denoting the interface between the two subdomains, and $f$ is assumed to be constant on $\\Omega_1$ and $\\Omega_2$, respectively. This tutorial shows how to create a nullspace, as required by pure Neumann boundary conditions on $\\partial\\Omega$, first in a case without domain decomposition and then in a case with domain decomposition (i.e., restrictions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9064cde1-da83-469b-82fe-dd40714fb275",
   "metadata": {},
   "outputs": [],
   "source": [
    "import typing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1effec2d-f2fb-45ab-8f21-63f984845cc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.fem\n",
    "import dolfinx.fem.petsc\n",
    "import dolfinx.io\n",
    "import dolfinx.la\n",
    "import gmsh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import petsc4py.PETSc\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bfa8650-9e4f-4bf2-bb68-e1d8f4caa040",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a11125b",
   "metadata": {},
   "source": [
    "### Geometrical parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea10c15c",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = 3\n",
    "mesh_size = 1. / 4."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0729d77",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "880377c5-8e37-474f-b2f7-6f0105d8e42d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.initialize()\n",
    "gmsh.model.add(\"mesh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b891ff83-5630-46b1-b042-c751ca33bc1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "p0 = gmsh.model.geo.addPoint(0.0, 0.0, 0.0, mesh_size)\n",
    "p1 = gmsh.model.geo.addPoint(0.0, +r, 0.0, mesh_size)\n",
    "p2 = gmsh.model.geo.addPoint(0.0, -r, 0.0, mesh_size)\n",
    "c0 = gmsh.model.geo.addCircleArc(p1, p0, p2)\n",
    "c1 = gmsh.model.geo.addCircleArc(p2, p0, p1)\n",
    "l0 = gmsh.model.geo.addLine(p2, p1)\n",
    "line_loop_left = gmsh.model.geo.addCurveLoop([c0, l0])\n",
    "line_loop_right = gmsh.model.geo.addCurveLoop([c1, -l0])\n",
    "semicircle_left = gmsh.model.geo.addPlaneSurface([line_loop_left])\n",
    "semicircle_right = gmsh.model.geo.addPlaneSurface([line_loop_right])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d6ed6e-bca3-4dee-83df-1afe82a720ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.model.geo.synchronize()\n",
    "gmsh.model.addPhysicalGroup(1, [c0], 1)\n",
    "gmsh.model.addPhysicalGroup(1, [c1], 2)\n",
    "gmsh.model.addPhysicalGroup(1, [l0], 3)\n",
    "gmsh.model.addPhysicalGroup(2, [semicircle_left], 1)\n",
    "gmsh.model.addPhysicalGroup(2, [semicircle_right], 2)\n",
    "gmsh.model.mesh.generate(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6028f7ce-7ab7-415d-ab03-36a73d90aa66",
   "metadata": {},
   "outputs": [],
   "source": [
    "partitioner = dolfinx.mesh.create_cell_partitioner(dolfinx.mesh.GhostMode.shared_facet)\n",
    "mesh, subdomains, boundaries_and_interfaces = dolfinx.io.gmshio.model_to_mesh(\n",
    "    gmsh.model, comm=mpi4py.MPI.COMM_WORLD, rank=0, gdim=2, partitioner=partitioner)\n",
    "gmsh.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cabe4d3-31b9-499a-9057-7f30bb26b2c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)\n",
    "mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03781a2a-93b5-4e60-bc27-f895d354d098",
   "metadata": {},
   "outputs": [],
   "source": [
    "cells_Omega1 = subdomains.indices[subdomains.values == 1]\n",
    "cells_Omega2 = subdomains.indices[subdomains.values == 2]\n",
    "facets_Gamma = boundaries_and_interfaces.indices[boundaries_and_interfaces.values == 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fc0d4ff-97a3-4ada-87f3-aa814ff3ea1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "integration_entities_on_Gamma = dolfinx.fem.compute_integration_domains(\n",
    "    dolfinx.fem.IntegralType.interior_facet, mesh.topology, facets_Gamma)\n",
    "integration_entities_on_Gamma_reshaped = integration_entities_on_Gamma.reshape(-1, 4)\n",
    "connected_cells_to_Gamma = integration_entities_on_Gamma_reshaped[:, [0, 2]]\n",
    "subdomain_ordering = (\n",
    "    subdomains.values[connected_cells_to_Gamma[:, 0]] < subdomains.values[connected_cells_to_Gamma[:, 1]])\n",
    "if len(subdomain_ordering) > 0 and any(subdomain_ordering):\n",
    "    integration_entities_on_Gamma_reshaped[subdomain_ordering] = integration_entities_on_Gamma_reshaped[\n",
    "        subdomain_ordering][:, [2, 3, 0, 1]]\n",
    "integration_entities_on_Gamma = integration_entities_on_Gamma_reshaped.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fc0c819-a45c-49d2-b202-775de234b6b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = ufl.Measure(\"dx\", subdomain_data=subdomains)\n",
    "ds = ufl.Measure(\"ds\", subdomain_data=boundaries_and_interfaces)\n",
    "dS = ufl.Measure(\"dS\", domain=mesh, subdomain_data=[(3, np.array(integration_entities_on_Gamma, dtype=np.int32))])\n",
    "dS = dS(3)  # restrict to the interface, which has facet ID equal to 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3bf0a99-a6af-4289-9d2b-2c4523dd6da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c5bdab-5bf1-4578-a173-56c3efe4518e",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, subdomains, \"subdomains\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c126c64-c5f0-4776-bf66-880717414b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, boundaries_and_interfaces, \"boundaries and interfaces\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "476854bd",
   "metadata": {},
   "source": [
    "### Without domain decomposition\n",
    "Note that this part does not actually require `multiphenicsx`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84959583",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create finite element space\n",
    "V = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1018feac-3cb9-48bc-8fcb-117ce0f2f5f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "u = ufl.TrialFunction(V)\n",
    "v = ufl.TestFunction(V)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e92794e-424f-44d0-b952-b9939e6b299b",
   "metadata": {},
   "source": [
    "Denote by $r$ the radius of the circular domain $\\Omega$.\n",
    "We next define the forcing term $f$ as\n",
    "$$f = \\begin{cases}1, & \\text{in }\\Omega_1,\\\\\n",
    "2, & \\text{in }\\Omega_2.\n",
    "\\end{cases}$$\n",
    "Assume that the boundary data $g$ is constant on $\\partial\\Omega$. Imposing the necessary condition for existence one can find the following formula for $g$:\n",
    "$$ g = - \\frac{3}{4} r, \\qquad \\text{on }\\partial\\Omega.$$\n",
    "The following cells verify the validity of the necessary condition, up to a tolerance proportional to the mesh size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de865a64-5d11-47f6-9ebd-f5ff8000a09a",
   "metadata": {},
   "outputs": [],
   "source": [
    "one = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(1))\n",
    "area_Omega1 = mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(one * dx(1))), op=mpi4py.MPI.SUM)\n",
    "area_Omega2 = mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(one * dx(2))), op=mpi4py.MPI.SUM)\n",
    "length_partial_Omega = mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(one * ds)), op=mpi4py.MPI.SUM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "935d2704-8b65-4176-980f-a3c1d76dacf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.isclose(area_Omega1, np.pi * r**2 / 2, atol=1e-1)\n",
    "assert np.isclose(area_Omega2, np.pi * r**2 / 2, atol=1e-1)\n",
    "assert np.isclose(length_partial_Omega, 2 * np.pi * r, atol=1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "203116fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = 1.0\n",
    "f2 = 2.0\n",
    "g = - 3.0 / 4.0 * r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d41345b8-21e2-42c1-b62c-10a547063a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.isclose(f1 * area_Omega1 + f2 * area_Omega2, - g * length_partial_Omega, atol=1e-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba4a4d37-99e6-44bf-bba1-aaa53677a26d",
   "metadata": {},
   "source": [
    "The weak formulation of the problem is therefore\n",
    "$$\n",
    "\\text{find }u \\in V(\\Omega) := H^1(\\Omega)\n",
    "$$\n",
    "s.t.\n",
    "$$\n",
    "\\int_{\\Omega} \\nabla u \\cdot \\nabla v dx =\n",
    "\\int_{\\Omega} f \\; v dx +\n",
    "\\int_{\\partial\\Omega} g \\; v ds.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82dd8d6f-aacf-4398-8046-06330ca67aab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem forms\n",
    "a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx\n",
    "f = ufl.inner(f1, v) * dx(1) + ufl.inner(f2, v) * dx(2) + ufl.inner(g, v) * ds\n",
    "a_cpp = dolfinx.fem.form(a)\n",
    "f_cpp = dolfinx.fem.form(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82cf2fa4-98b2-4dbb-8b30-2367e274800b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the discrete system\n",
    "A = dolfinx.fem.petsc.assemble_matrix(a_cpp)\n",
    "A.assemble()\n",
    "F = dolfinx.fem.petsc.assemble_vector(f_cpp)\n",
    "F.ghostUpdate(addv=petsc4py.PETSc.InsertMode.ADD, mode=petsc4py.PETSc.ScatterMode.REVERSE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd4c320-605f-4359-86fc-62b0a06628b5",
   "metadata": {},
   "source": [
    "However, the formulation above still needs to account for the fact that the solution is determined up to a constant. In other words, the matrix `A` has a null space associated to the finite element representation of every constant function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d64236d3-eb66-44d6-b12a-8e92ecb48850",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = dolfinx.fem.Function(V)\n",
    "c.interpolate(lambda x: np.ones(x.shape[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f054dcf8-0030-46fe-a108-7fcc82357ac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = c.x.petsc_vec\n",
    "assert np.allclose(C.array, 1.0)  # because V is a Lagrange space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcd9247e-0909-4798-b501-5287a5a40b45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MatNullSpaceCreate expects the vectors to be orthonormal, which in this case simply\n",
    "# means that we should normalize the vector C\n",
    "C.scale(1 / C.norm())\n",
    "assert np.isclose(C.norm(), 1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba90f22-05e2-4ed4-bc65-283317a62394",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the PETSc nullspace vector and check that it is a valid nullspace of A\n",
    "nullspace = petsc4py.PETSc.NullSpace().create(vectors=[C], comm=mesh.comm)\n",
    "assert nullspace.test(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fe08d9d-6ecb-4566-860b-000b51903076",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For convenience, we explicitly inform PETSc that A is symmetric, so that it automatically\n",
    "# sets the nullspace of A^T too (see the documentation of MatSetNullSpace).\n",
    "A.setOption(petsc4py.PETSc.Mat.Option.SYMMETRIC, True)\n",
    "A.setOption(petsc4py.PETSc.Mat.Option.SYMMETRY_ETERNAL, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a33188-5cf3-4378-ac14-56ad26508237",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the nullspace\n",
    "A.setNullSpace(nullspace)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d566d685-3f7e-4207-8e8f-a266de7d2b28",
   "metadata": {},
   "source": [
    "The documentation of `PETSc` suggests to orthogonalize `F` to the null space of `A^T`. We note that this is theoretically unnecessary here, because $g$ has been chosen to satisfy the compatibility condition with `f`. Still, we carry out the orthogonalization anyway because, in practice, the finite element vector `F` actually has a component in the null space of `A^T` e.g. because of the fact that the mesh does not represent perfectly a circle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58cfd0ed-5f31-462c-a9f1-1c5003251c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Orthogonalize F to the null space of A^T\n",
    "nullspace.remove(F)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71ac53c0-c513-42d7-aef3-dcaa22245c08",
   "metadata": {},
   "source": [
    "We next configure a direct solver with `MUMPS` to solve the linear system.\n",
    "Note that `MUMPS` requires to explicitly set two options in the case of singular linear systems. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b1a5a91-d7d2-4957-ab28-99684d488b72",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = dolfinx.fem.Function(V)\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.getPC().setFactorSetUpSolverType()\n",
    "ksp.getPC().getFactorMatrix().setMumpsIcntl(icntl=24, ival=1)  # detect null pivots\n",
    "ksp.getPC().getFactorMatrix().setMumpsIcntl(icntl=25, ival=0)  # do not compute null space again\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, solution.x.petsc_vec)\n",
    "solution.x.petsc_vec.ghostUpdate(\n",
    "    addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4b69fc2-24de-4e76-8acc-1bd012c9062b",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(solution, \"u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c5387bc-065d-47d1-bf78-7d4c2e69c0ab",
   "metadata": {},
   "source": [
    "We finally note that setting the null space is not equivalent to imposing a constraint\n",
    "$$\\int_\\Omega u \\; dx = 0$$\n",
    "like in the [\"Singular Poisson\" demo](https://fenicsproject.org/olddocs/dolfin/dev/python/demos/singular-poisson/demo_singular-poisson.py.html) in legacy FEniCS. In other words, the linear solver will fix the undetermined constant in a way that the solution does not necessarily have zero average. If interested in enforcing the zero average constraint, the user can postprocess the obtained solution simply subtracting the average of the computed solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35c251d6-70d1-407c-ac93-133c1dcfcb0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_average(  # type: ignore[no-any-unimported]\n",
    "    u: typing.Union[dolfinx.fem.Function, tuple[dolfinx.fem.Function, dolfinx.fem.Function]]\n",
    ") -> petsc4py.PETSc.ScalarType:\n",
    "    \"\"\"Compute average of the solution.\"\"\"\n",
    "    if not isinstance(u, tuple):\n",
    "        u = (u, u)\n",
    "    else:\n",
    "        assert len(u) == 2\n",
    "    return mesh.comm.allreduce(\n",
    "        dolfinx.fem.assemble_scalar(dolfinx.fem.form(u[0] * dx(1) + u[1] * dx(2))), op=mpi4py.MPI.SUM\n",
    "    ) / (area_Omega1 + area_Omega2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "819016b7-28a9-4611-804a-d1aa1c1a7d90",
   "metadata": {},
   "outputs": [],
   "source": [
    "def subtract_average(  # type: ignore[no-any-unimported]\n",
    "    u: dolfinx.fem.Function, average_u: petsc4py.PETSc.ScalarType,\n",
    "    active_dofs: typing.Optional[np.typing.NDArray[np.int32]] = None\n",
    ") -> None:\n",
    "    \"\"\"Post-process the solution so that it has zero average.\"\"\"\n",
    "    with u.x.petsc_vec.localForm() as u_local:\n",
    "        if active_dofs is None:\n",
    "            u_local[:] -= average_u\n",
    "        else:\n",
    "            u_local[active_dofs] -= average_u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87393b99-2108-4059-869e-9e1651bfe207",
   "metadata": {},
   "outputs": [],
   "source": [
    "average_u = compute_average(solution)\n",
    "assert np.isclose(average_u, 0.11117, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24983809-de1c-4fe6-8842-c2f6a720c316",
   "metadata": {},
   "outputs": [],
   "source": [
    "subtract_average(solution, average_u)\n",
    "assert np.isclose(compute_average(solution), 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4aed78e5-02d4-4737-b55e-4884781481dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "del average_u"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e199b6d7-78e5-4272-8829-2699fac4c87f",
   "metadata": {},
   "source": [
    "### With domain decomposition\n",
    "We next perform a domain decomposition of $\\Omega$ as $\\Omega_1 \\cup \\Omega_2$. Similarly to [the interface example in tutorial 03](../03_lagrange_multipliers/tutorial_lagrange_multipliers_interface.ipynb),\n",
    "we need to introduce a lagrange multiplier to handle the continuity of the solution across\n",
    "the interface $\\Gamma$ between $\\Omega_1$ and $\\Omega_2$.\n",
    "\n",
    "The resulting weak formulation is:\n",
    "$$\n",
    "\\text{find }u_1 \\in V(\\Omega_1), u_2 \\in V(\\Omega_2), \\eta \\in E(\\Gamma) \\subset L^2(\\Gamma)\n",
    "$$\n",
    "s.t.\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\int_{\\Omega_1} \\nabla u_1 \\cdot \\nabla v_1 dx + \\int_{\\Gamma} \\lambda \\; v_1  ds \n",
    "= \\int_{\\Omega_1} f \\; v_1 dx + \\int_{\\partial\\Omega_1 \\setminus \\Gamma} g \\; v_1 ds,\n",
    "& \\qquad \\forall v_1 \\in V(\\Omega_1)\\\\\n",
    "\\int_{\\Omega_2} \\nabla u_2 \\cdot \\nabla v_2 dx - \\int_{\\Gamma} \\lambda \\; v_2 ds \n",
    "= \\int_{\\Omega_2} f \\; v_2 dx + \\int_{\\partial\\Omega_2 \\setminus \\Gamma} g \\; v_2 ds,\n",
    "& \\qquad \\forall v_2 \\in V(\\Omega_2)\\\\\n",
    "\\int_{\\Gamma} \\eta \\; (u_1 - u_2) ds = 0,\n",
    "& \\qquad \\forall \\eta \\in E(\\Gamma).\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Also in this case the solution $u_1$ and $u_2$ are defined up to a constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98d0ad3-e98a-4e9e-b495-b807fb5b2fac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function spaces\n",
    "V1 = V.clone()\n",
    "V2 = V.clone()\n",
    "M = V.clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ee1d65-0443-4891-bfe3-a2c1da2cf79c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define restrictions\n",
    "dofs_V1_Omega1 = dolfinx.fem.locate_dofs_topological(V1, subdomains.dim, cells_Omega1)\n",
    "dofs_V2_Omega2 = dolfinx.fem.locate_dofs_topological(V2, subdomains.dim, cells_Omega2)\n",
    "dofs_M_Gamma = dolfinx.fem.locate_dofs_topological(M, boundaries_and_interfaces.dim, facets_Gamma)\n",
    "restriction_V1_Omega1 = multiphenicsx.fem.DofMapRestriction(V1.dofmap, dofs_V1_Omega1)\n",
    "restriction_V2_Omega2 = multiphenicsx.fem.DofMapRestriction(V2.dofmap, dofs_V2_Omega2)\n",
    "restriction_M_Gamma = multiphenicsx.fem.DofMapRestriction(M.dofmap, dofs_M_Gamma)\n",
    "restriction = [restriction_V1_Omega1, restriction_V2_Omega2, restriction_M_Gamma]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe16d6e8-0fc2-462f-b1a4-4a5298808590",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "(u1, u2, l) = (ufl.TrialFunction(V1), ufl.TrialFunction(V2), ufl.TrialFunction(M))\n",
    "(v1, v2, m) = (ufl.TestFunction(V1), ufl.TestFunction(V2), ufl.TestFunction(M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dbe75d3-4186-47fc-86ff-1c9230fa894a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem block forms\n",
    "zero = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0))\n",
    "a_dd = [[ufl.inner(ufl.grad(u1), ufl.grad(v1)) * dx(1), None, ufl.inner(l(\"-\"), v1(\"-\")) * dS],\n",
    "        [None, ufl.inner(ufl.grad(u2), ufl.grad(v2)) * dx(2), - ufl.inner(l(\"+\"), v2(\"+\")) * dS],\n",
    "        [ufl.inner(u1(\"-\"), m(\"-\")) * dS, - ufl.inner(u2(\"+\"), m(\"+\")) * dS, None]]\n",
    "f_dd = [ufl.inner(f1, v1) * dx(1) + ufl.inner(g, v1) * ds(1),\n",
    "        ufl.inner(f2, v2) * dx(2) + ufl.inner(g, v2) * ds(2),\n",
    "        ufl.inner(zero, m(\"-\")) * dS]\n",
    "a_dd_cpp = dolfinx.fem.form(a_dd)\n",
    "f_dd_cpp = dolfinx.fem.form(f_dd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bde17499-f7cc-4536-9f02-8a6d7926366e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system\n",
    "A_dd = multiphenicsx.fem.petsc.assemble_matrix_block(a_dd_cpp, restriction=(restriction, restriction))\n",
    "A_dd.assemble()\n",
    "F_dd = multiphenicsx.fem.petsc.assemble_vector_block(f_dd_cpp, a_dd_cpp, restriction=restriction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b24bee-c97c-4526-9ecc-3cfde7e72f2c",
   "metadata": {},
   "source": [
    "We then set a null space which specifies that $u_1$ and $u_2$ are determined up to a constant $c_1$ and $c_2$, respectively. The null space will consist of a block vector `C_dd`, which entries are grouped in the following way: first, DOFs associated to $V_1$, then DOFs associated to $V_2$, finally DOFs associated to $M$. The block vector is initialized to zero, and its non-zero entries are as constructed follows:\n",
    "1. copy the values of `C` associated to every DOF of $V$ in $\\Omega_1$ into the corresponding DOFs of $V_1$ in the first block of `C_dd`;\n",
    "2. copy the values of `C` associated to every DOF of $V$ in $\\Omega_2$ into the corresponding DOFs of $V_2$ in the second block of `C_dd`;\n",
    "3. leave entries in the third block of `C_dd` equal to zero.\n",
    "\n",
    "This does indeed construct a null space for the block system because:\n",
    "* with 1 and 3, we represent that the bilinear form $\\int_{\\Omega_1} \\nabla u_1 \\cdot \\nabla v_1 \\; dx + \\int_{\\Gamma} \\lambda \\; v_1 \\; ds$ on the left-hand side of the first equation of the system has a null space when $u_1$ assumes any constant value $c_1$ and $\\lambda$ is equal to zero;\n",
    "* with 2 and 3, we represent that the bilinear form $\\int_{\\Omega_2} \\nabla u_2 \\cdot \\nabla v_2 \\; dx + \\int_{\\Gamma} \\lambda \\; v_2 \\; ds$ on the left-hand side of the second equation of the system has a null space when $u_2$ assumes any constant value $c_2$ and $\\lambda$ is equal to zero;\n",
    "* since `C` was a vector composed of all entries with the same value, the non-zero entries of `C_dd` as a result of 1 have the same value of the non-zero entries filled as a result of 2. This implies that $c_1$ is actually equal to $c_2$, and therefore the left-hand side of the third equation $\\int_{\\Gamma} \\eta \\; (u_1 - u_2) ds$ is equal to zero when $u_1 = c_1 = c_2 = u_2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c02d2fd-a332-4ee9-919a-87f89c52d847",
   "metadata": {},
   "outputs": [],
   "source": [
    "C_dd = multiphenicsx.fem.petsc.create_vector_block(f_dd_cpp, restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c79ba44-afef-4c9e-b501-a4097c8acf08",
   "metadata": {},
   "outputs": [],
   "source": [
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(\n",
    "        C_dd, [V1.dofmap, V2.dofmap, M.dofmap], restriction) as C_dd_wrapper:\n",
    "    for C_dd_component_local, data_vector in zip(C_dd_wrapper, (C, C, None)):\n",
    "        if data_vector is not None:  # skip third block\n",
    "            with data_vector.localForm() as data_vector_local:\n",
    "                C_dd_component_local[:] = data_vector_local"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b26ba050-4e38-4d73-a327-55475398487c",
   "metadata": {},
   "source": [
    "Then, create a nullspace using a similar code as for the case without domain decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697bc2e0-9d6d-425c-99bd-9b6d2f099237",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MatNullSpaceCreate expects the vectors to be orthonormal\n",
    "C_dd.scale(1 / C_dd.norm())\n",
    "assert np.isclose(C_dd.norm(), 1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b8f0733-abc9-4917-94c1-fb591f973fbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the PETSc nullspace vector and check that it is a valid nullspace of A_dd\n",
    "nullspace_dd = petsc4py.PETSc.NullSpace().create(vectors=[C_dd], comm=mesh.comm)\n",
    "assert nullspace_dd.test(A_dd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18e1189-71a3-46a9-b070-a2d78c63f60a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inform PETSc that A_dd is symmetric\n",
    "A_dd.setOption(petsc4py.PETSc.Mat.Option.SYMMETRIC, True)\n",
    "A_dd.setOption(petsc4py.PETSc.Mat.Option.SYMMETRY_ETERNAL, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50f45e2e-73c9-4d62-8d20-4eeb4d65e5f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the nullspace\n",
    "A_dd.setNullSpace(nullspace_dd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e540d154-be2a-4cf5-b75b-0dd09bb0b683",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Orthogonalize F_dd to the null space of A_dd^T\n",
    "nullspace_dd.remove(F_dd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d60de8-4e23-477a-a2d1-8826b9494f98",
   "metadata": {},
   "source": [
    "Finally, solve the system with the same solver as in the case without domain decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acc89273-cbcf-4e09-85da-0322b227d402",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "u1u2l_dd = multiphenicsx.fem.petsc.create_vector_block(f_dd_cpp, restriction=restriction)\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A_dd)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.getPC().setFactorSetUpSolverType()\n",
    "ksp.getPC().getFactorMatrix().setMumpsIcntl(icntl=24, ival=1)\n",
    "ksp.getPC().getFactorMatrix().setMumpsIcntl(icntl=25, ival=0)\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F_dd, u1u2l_dd)\n",
    "u1u2l_dd.ghostUpdate(\n",
    "    addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "630aa2eb-69cf-4e06-84c5-a52210d60cc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "(u1_dd, u2_dd, l_dd) = (dolfinx.fem.Function(V1), dolfinx.fem.Function(V2), dolfinx.fem.Function(M))\n",
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(\n",
    "        u1u2l_dd, [V1.dofmap, V2.dofmap, M.dofmap], restriction) as u1u2l_wrapper:\n",
    "    for u1u2l_wrapper_local, component in zip(u1u2l_wrapper, (u1_dd, u2_dd, l_dd)):\n",
    "        with component.x.petsc_vec.localForm() as component_local:\n",
    "            component_local[:] = u1u2l_wrapper_local\n",
    "u1u2l_dd.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed15e9fc-7fce-471f-afd0-87c27ac053e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u1_dd, \"u1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca6b2e0c-98fb-4697-b95d-30ce24f880c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u2_dd, \"u2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00c280ee-6bb2-4027-8b97-dd8033735976",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(l_dd, \"l\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "970dff07-6c50-412f-a8c8-0ef7daf2cd07",
   "metadata": {},
   "source": [
    "In order to make a comparison to $u$ in the final section of this notebook, we subtract the average of the sum of the solutions $u_1$ and $u_2$, since there is no guarantee that the solutions without and with domain decomposition set the same constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6bc79c4-d0c6-4d06-8c39-f0c6e3fe62cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "average_u_dd = compute_average((u1_dd, u2_dd))\n",
    "assert np.isclose(average_u_dd, 0.09797, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52f379e9-48a7-498d-9eec-632b5693bc03",
   "metadata": {},
   "outputs": [],
   "source": [
    "subtract_average(u1_dd, average_u_dd, dofs_V1_Omega1)\n",
    "subtract_average(u2_dd, average_u_dd, dofs_V2_Omega2)\n",
    "assert np.isclose(compute_average((u1_dd, u2_dd)), 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ce3a928-736c-46ea-a019-ee1b856e542a",
   "metadata": {},
   "outputs": [],
   "source": [
    "del average_u_dd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9eff1742-c5f7-4aec-96d1-9040d7dfaac0",
   "metadata": {},
   "source": [
    "### Error computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ada68804-44bd-4fe9-9b8f-eaaa8cfdc338",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_norm_Omega1 = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(solution, solution) * dx(1))),\n",
    "    op=mpi4py.MPI.SUM))\n",
    "u_norm_Omega2 = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(solution, solution) * dx(2))),\n",
    "    op=mpi4py.MPI.SUM))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc0c046-9694-4ff0-b561-30df190a63a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "err1_norm_dd = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(solution - u1_dd, solution - u1_dd) * dx(1))),\n",
    "    op=mpi4py.MPI.SUM))\n",
    "err2_norm_dd = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(solution - u2_dd, solution - u2_dd) * dx(2))),\n",
    "    op=mpi4py.MPI.SUM))\n",
    "print(\"Relative error on subdomain 1\", err1_norm_dd / u_norm_Omega1)\n",
    "print(\"Relative error on subdomain 2\", err2_norm_dd / u_norm_Omega2)\n",
    "assert np.isclose(err1_norm_dd / u_norm_Omega1, 0., atol=3.e-5)\n",
    "assert np.isclose(err2_norm_dd / u_norm_Omega2, 0., atol=3.e-5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
